# coding: UTF-8
import torch
import sys
from config import *
import numpy as np
import time
from utils.utils import *
from model.TextCNN import Config,Model

def train(cfg,model,train_iter, dev_iter, test_iter):
    start_time = time.time()
    model.train()
    optimizer = torch.optim.Adam(model.parameters(), lr=cfg.learning_rate)
    total_batch = 0  # 记录进行到多少batch
    dev_best_loss = float('inf')
    last_improve = 0  # 记录上次验证集loss下降的batch数
    flag = False  # 记录是否很久没有效果提升
    for epoch in range(cfg.num_epochs):
        print('Epoch [{}/{}]'.format(epoch + 1, cfg.num_epochs))
        for i,(trains,label) in enumerate(train_iter):
            print(label.shape,type(label))
            print(type(trains),trains.shape)
            sys.exit(0)

def main():
    wiki_w2c =  pre_train_dir / "embedding_wiki_w2c.npz"
    wiki_cw2c =  pre_train_dir / "embedding_wiki_cw2c.npz"
    print(wiki_cw2c)
    print(os.path.exists(wiki_cw2c))
    random_seed = 1
    np.random.seed(1)
    torch.manual_seed(1)
    torch.cuda.manual_seed_all(1)
    torch.backends.cudnn.deterministic = True  # 保证每次结果一样
    cfg = Config(embedding=wiki_cw2c)

    start_time = time.time()
    print("Loading data...")
    sougo_dator = SougouDataset(sougouCS_dir=sougouCS_dir)
    train_set,dev_set,test_set,word2index_dict,index2word_dict = sougo_dator.getDataset()
    # print("训练集尺寸：{}, 标签尺寸：{}".format(train_set[0].shape, train_set[1].shape))
    # print("验证集尺寸：{}, 标签尺寸：{}".format(dev_set[0].shape, dev_set[1].shape))
    # print("测试集尺寸：{}, 标签尺寸：{}".format(test_set[0].shape, test_set[1].shape))
    # print("词表数量：", len(word2index_dict.keys()), len(index2word_dict))

    train_iter = build_iterator(train_set, cfg)
    dev_iter = build_iterator(dev_set, cfg)
    test_iter = build_iterator(test_set, cfg)
    time_dif = get_time_dif(start_time)

    print("Time usage:", time_dif)
    cfg.n_vocab = len(word2index_dict)


    model = Model(config=cfg).to(device=cfg.device)
    print(model)
    train(cfg, model, train_iter, dev_iter, test_iter)
    pass
if __name__ == '__main__':
    main()
